package simplytypedinference


/**
 * unification algorithm from TAPL chapter 22
 */
object Unification {
  import inference._
  
  def assertEquals(s:String,actual:Any,expected:Any) {
    if (actual == expected) println (s+":OK  == "+actual) 
    else println(s+ ":ERROR     "+actual +" != " +expected)
  }
  
  def main(args: Array[String]) {
    println(unify(List((TyVar("x"),TyVar("x")))))
    val pt1 = unify(List( TyVar("X") -> TyInt,
        TyVar("Y") -> TyArr(TyVar("X"),TyVar("X"))))
    assertEquals("ex22_4_3_pt1",pt1, List(Subst(TyVar("X"), TyInt) , Subst(TyVar("Y") , TyArr(TyInt, TyInt)) ))
        
    
    val pt2 = unify(List( TyArr(TyInt,TyInt) -> TyArr(TyVar("X"),TyVar("Y")))) 
    assertEquals("ex22_4_3_pt2",pt2, List(Subst(TyVar("X"), TyInt) , 
                                          Subst(TyVar("Y"), TyInt) ))

    val pt3 = unify(List( TyArr(TyVar("X"),TyVar("Y")) 
                             -> TyArr(TyVar("Y"),TyVar("Z")), 
                          TyVar("Z") 
                             -> TyArr(TyVar("U"),TyVar("W"))))
    assertEquals("ex22_4_3_pt3",pt3, List(Subst(TyVar("X"), TyArr(TyVar("U"),TyVar("W"))),
                                          Subst(TyVar("Y"), TyArr(TyVar("U"),TyVar("W"))),
                                          Subst(TyVar("Z"), TyArr(TyVar("U"),TyVar("W")))))
  
    
    try {
    println("ex22_4_3_pt4: "+unify(List( TyInt      -> TyArr(TyInt, TyVar("Y"))))) 
    }catch { case e:Exception => println( "ex22_4_3_pt4 "+e.getMessage())
    }
    try {
    println("ex22_4_3_pt5: "+unify(List( TyVar("Y") -> TyArr(TyInt, TyVar("Y")))))
    }catch { case e:Exception => println( "ex22_4_3_pt5 "+e.getMessage())
    }
    
    val pt6 = unify(Nil)
    assertEquals("ex22_4_3_pt6",pt6, Nil)
        
  }
}

object inference {

type Constraint = (LType, LType)



def unify(constraints: List[Constraint]) : List[Subst] = constraints match{
  case Nil => Nil
  case (s,t)::cs => 
  	if ( s == t) unify(cs)
        else (s,t) match {
          case (s_ @TyVar(n), t_) if (!(freevars(t) exists {x => x==n})) =>
            Subst(s_ ,t_ ) o unify(Subst(s_ ,t_)(cs) )
          case (s_, t_ @TyVar(n) )if (!(freevars(s) exists {x => x==n})) =>
            Subst(t_, s_) o unify(Subst(t_, s_)(cs) )
          case (s_ @ TyArr(s1,s2), t_ @ TyArr(t1,t2))                    =>
            unify((s1,t1) :: (s2,t2):: cs)
          case (TyBool(), TyBool())                                      =>
            unify(cs)
          case (TyInt(),  TyInt() )                                      =>
            unify(cs)
          case (TyUnit(), TyUnit())                                      =>
            unify(cs)
          case _ => throw new TypeExc("unification failed : "+s +" not unifiable with " +t)
        }
  	
}

def freevars(t: LType) : List[String]= t match {
  case TyVar(n)   => List(n)
  case TyUnit()   => Nil
  case TyInt()    => Nil
  case TyBool()   => Nil
  case TyArr(t1,t2)=> freevars(t1):::freevars(t2)
}

  
case class Subst(v: TyVar, t: LType) extends Product2[TyVar,LType]{
  
  
  
  override def toString = "[" +v +" -> " +t +"]"
    
  def apply(cons: List[Constraint]) : List[Constraint] = maptupleList(cons, subst)
  def apply2(cons: List[Subst]) : List[Subst] = mapSubstList(cons, substprecise,subst)
  
  def _1 = v
  def _2 = t
  
  def o(xs:List[Subst]) = {
    var fullpart = this
    xs.foreach(x => {fullpart = x.apply2( List(fullpart)) . head })
    fullpart :: xs
  }
  //this:: mapSubstList(xs, substprecise,subst)
  
  def substprecise(s:TyVar) :TyVar = s match {
  case v_ @ TyVar(_) if v == v_    => t. asInstanceOf [TyVar]
  case v_ @ TyVar(_)               => v_
  }
  def subst(s:LType):LType = s match {
  case v_ @ TyVar(_) if v == v_    => t
  case v_ @ TyVar(_)               => v_
  case TyArr(s1,s2)                => TyArr(subst(s1),subst(s2))
  
  case TyUnit()                    => s
  case TyInt()                     => s
  case TyBool()                    => s
  
    }
  //todo: fix all this: subst is the whole chain, not just one single substitution
  def maptupleList[A,B](l: List[(A,A)], f:A => B) :List[(B,B)]    = 
        l. flatMap( {case (ts1,ts2)      => List((f(ts1), f(ts2))) })
  def mapSubstList(l: List[Subst], f1: TyVar => TyVar,f2: LType=>LType) :List[Subst] = 
        l. flatMap( {case Subst(ts1,ts2) => List(Subst(f1(ts1),f2(ts2)))})
}

}